# -*- encoding: utf-8 -*-
"""
    @File      : Aini
    @Time      : 2023-08-21 20:27:48
    @Function  : 加密解密目录

    加密：使用 ecb 分组加密模式加密指定的目录
        要加密的目录必须指定而且是指定绝对路径
        加密后保存的目录可以不指定，如果不指定会保存要加密目录的同级目录下
        如果 key 没有指定会生成 32 位的安全 key
        目录下的文件加密时会生成 base64 编码的 解密密码，密钥对在 .vital 目录下
    解密：要解密的目录必须指定而且必须是绝对路径
        解密后存储的目录可以不指定，会保存到同级目录下
        必须指定密码，必须指定私钥，私钥用于解密密码，密码用于解密数据
"""

import os
import base64
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto.Util.Padding import pad
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Util.Padding import unpad
from Crypto.Random import get_random_bytes


class myEncryptAndDecrypt:
    """我的 加密 解密

    ecb_key: 对称加密密钥长度，密钥长度只能是 16 24 32 其中之一
    可用方法:
        encrypt: 加密目录
        decrypt: 解密目录
    """

    # 白名单，可以是文件或文件夹，不加密
    white_list = ['.vital']

    def __init__(self) -> None:
        self.ecb_key = get_random_bytes(32)

    def _handle_output_directory(self, type: str, input_directory: str) -> str:
        """加密或解密后导出的目录路径
        type: 当前是加密还是解密
            - encrypt 是加密
            - decrypt 是解密
        input_directory: 要加密的目录
            会在目录名后面添加 _encrypted 或 _decrypted
        return: 加密目录的路径
        """
        if type == 'encrypt':
            output_directory = os.path.abspath(os.path.join(input_directory, "..", input_directory+"_encrypted"))
        if type == 'decrypt':
            output_directory = os.path.abspath(os.path.join(input_directory, "..", input_directory+"_decrypted"))
        if not os.path.exists(output_directory):
            os.makedirs(output_directory)
        return output_directory

    def _encrypt_file(self, file_path: str, output_path: str) -> None:
        """文件加密, 存到指定路径
        file_path: 加密的文件路径
        output_path: 加密后存储的文件路径
        """
        # 读取文件,加密文件内容
        with open(file_path, 'rb') as f_in:
            # ECB 加密对象
            cipher = AES.new(self.ecb_key, AES.MODE_ECB)
            encrypted_data = cipher.encrypt(pad(f_in.read(), AES.block_size))
        # 写入加密后的文件
        with open(output_path, 'wb') as f_out:
            f_out.write(encrypted_data)

    def _encoding_file_name(self, file_name: str) -> str:
        """文件名自定义编码
        file_name: 编码的文件名
        return: 编码后的文件名
        """
        key = file_name[::-1]
        lists = []
        for i in key:
            utfnum = str(ord(i))
            lens = len(utfnum)
            lens = str(lens)
            lists.append(lens)
            lists.append(utfnum)
        isok = ''.join(lists)
        isok = int(isok)
        return hex(isok)[2:]

    def _handle_encrypt_directory(self, input_directory: str, output_directory: str) -> None:
        """递归遍历给定的目录，加密目录下的文件
        input_directory: 要加密的目录
        output_directory: 加密后存储的目录
        """
        for name in os.listdir(input_directory):
            # 跳过白名单内的文件或目录
            if name in self.white_list:
                continue
            file_path = os.path.join(input_directory, name)
            # 处理子目录
            if os.path.isdir(file_path):
                output_path = os.path.join(output_directory, name)
                # 创建文件夹
                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                self._handle_encrypt_directory(file_path, output_path)
            # 处理文件
            else:
                # 文件名编码
                encoded_name = self._encoding_file_name(name)
                output_path = os.path.join(output_directory, encoded_name)
                self._encrypt_file(file_path, output_path)

    def _decrypt_file(self, file_path: str, output_path: str) -> None:
        """ 文件加密, 存到指定路径
        file_path: 要解密的文件路径
        output_path: 解密后存储的文件路径
        """
        # 读取文件,解密文件内容
        with open(file_path, 'rb') as f_in:
            # ECB 对象
            cipher = AES.new(self.ecb_key, AES.MODE_ECB)
            decrypted_data = unpad(cipher.decrypt(f_in.read()), AES.block_size)
        # 写入解密后的文件
        with open(output_path, 'wb') as f_out:
            f_out.write(decrypted_data)

    def _decoding_file_name(self, file_name: str) -> str:
        """文件名解码
        file_name: 解码的文件名
        return: 解码后的文件名
        """
        key = str(int('0x'+file_name, 16))
        temp = 0
        list = []
        for a in key:
            index = int(key[temp])
            start = temp+1
            end = temp+index+1
            num = key[start:end]
            list.append(num)
            temp = end
            if(temp >= len(key)):
                break
        list = list[::-1]
        oklist = []
        for i in list:
            i = int(i)
            isok = chr(i)
            oklist.append(isok)
        return ''.join(oklist)

    def _handle_decrypt_directory(self, input_directory: str, output_directory: str) -> None:
        """递归遍历给定的目录，加密目录下的文件
        input_directory: 要加密的目录
        output_directory: 加密后存储的目录
        """
        for name in os.listdir(input_directory):
            # 跳过白名单内的文件或目录
            if name in self.white_list:
                continue
            file_path = os.path.join(input_directory, name)
            # 处理子目录
            if os.path.isdir(file_path):
                output_path = os.path.join(output_directory, name)
                # 创建文件夹
                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                self._handle_decrypt_directory(file_path, output_path)
            # 处理文件
            else:
                # 文件名解码
                decoded_name = self._decoding_file_name(name)
                output_path = os.path.join(output_directory, decoded_name)
                self._decrypt_file(file_path, output_path)

    def _encrypt_ecb_key(self, key_pub_path: str, output_directory: str) -> None:
        """ 使用公钥加密非对称加密用的 key
        key_pub_path: 公钥文件路径
        output_directory: 加密后的文件保存路径
        """
        # 生成保存密钥和密码的目录
        vital_path = os.path.join(output_directory, '.vital')
        if not os.path.exists(vital_path):
            os.makedirs(vital_path)
            print('已创建存储密钥和密码的目录 ', vital_path)
        # 读取公钥
        with open(key_pub_path, 'rb') as f_in:
            key_pub = RSA.import_key(f_in.read())
            cipher = PKCS1_OAEP.new(key_pub)
            ciphertext = cipher.encrypt(self.ecb_key)
        # 保存使用公钥加密后的 key
        cipher_path = os.path.join(vital_path, 'cipher_base64')
        with open(cipher_path, 'wb') as f_out:
            f_out.write(base64.b64encode(ciphertext))
        print('key 已加密, 密文已编码 ', cipher_path)

    def _decrypt_ecb_key(self, key_path: str, cipher_base64_path: str) -> bytes:
        """ 解密非对称加密用的 key 
        key_path: 私钥文件
        cipher_base64: 密码文件
        return: 解密后的key
        """
        # 私钥读取
        with open(key_path, 'rb') as f:
            key = RSA.import_key(f.read())
        # 密码读取
        with open(cipher_base64_path, 'rb') as f:
            ciphertext = base64.b64decode(f.read())
        cipher = PKCS1_OAEP.new(key)
        return cipher.decrypt(ciphertext)

    def _generate_private_key(self, output_directory: str) -> str:
        """ 生成 RSA 密钥对 
        output_directory: 密钥对保存的路径
        return: 公钥路径
        """
        private_key = RSA.generate(2048)
        # 生成保存密钥和密码的目录
        vital_path = os.path.join(output_directory, '.vital')
        if not os.path.exists(vital_path):
            os.makedirs(vital_path)
            print('已创建存储密钥和密码的目录 ', vital_path)
        # 私钥路径
        key_path = os.path.join(vital_path, 'key')
        with open(key_path, 'wb') as f:
            f.write(private_key.export_key('PEM'))
        # 公钥路径
        key_pub_path = os.path.join(vital_path, 'key.pub')
        with open(key_pub_path, 'wb') as f:
            f.write(private_key.public_key().export_key('PEM'))
        print('密钥已生成 ', key_path)
        print('私钥已生成 ', key_pub_path)
        return key_pub_path

    def encrypt(self, input_directory: str, output_directory: str = None, public_key_path: str = None) -> None:
        """加密
        input_directory: 要加密的目录
            必须指定绝对路径
        output_directory: 加密后保存的目录
            可选，如果不指定会保存到要加密文件夹的同级目录下
        public_key_path: 加密key用的公钥路径
            可选，如果不指定，自动生成一个密钥对，密钥保存到加密后目录下
        """
        if not os.path.isdir(input_directory):
            raise FileNotFoundError('要加密的目录不存在')
        if not output_directory:
            output_directory = self._handle_output_directory('encrypt', input_directory)
        if not os.path.isdir(output_directory):
            raise FileNotFoundError('加密后存储的目录不存在')
        if public_key_path and not os.path.isfile(public_key_path):
            raise FileNotFoundError('公钥文件不存在')
        if not public_key_path:
            public_key_path = self._generate_private_key(output_directory)
        self._encrypt_ecb_key(public_key_path, output_directory)
        self._handle_encrypt_directory(input_directory, output_directory)

    def decrypt(self, input_directory: str, output_directory: str = None, key_path: str = None, cipher_base64_path: str = None) -> None:
        """解密
        input_directory: 要解密的目录
            必须指定绝对路径
        output_directory: 解密后保存的目录
            可选，如果不指定会保存到要解密文件夹的同级目录下
        key_path: 解密用的私钥
            必须指定绝对路径
        cipher_base64_path: ecb_key 密码
            必须指定绝对路径
        """
        if not key_path:
            raise FileNotFoundError('必须指定密钥')
        if not os.path.isfile(key_path):
            raise FileNotFoundError('私钥不存在')
        if not cipher_base64_path:
            raise FileNotFoundError('必须指定密码')
        if not os.path.isfile(cipher_base64_path):
            raise FileNotFoundError('密码文件不存在')
        if not key_path or not os.path.isfile(key_path):
            raise FileNotFoundError('私钥不存在')
        if not os.path.isdir(input_directory):
            raise FileNotFoundError('要解密的目录不存在')
        if not output_directory:
            output_directory = self._handle_output_directory('decrypt', input_directory)
        if not os.path.isdir(output_directory):
            raise FileNotFoundError('解密后存储的目录不存在')
        self.ecb_key = self._decrypt_ecb_key(key_path, cipher_base64_path)
        self._handle_decrypt_directory(input_directory, output_directory)


